from math import log
import operator

# https://www.cnblogs.com/ku1274755259/p/11108940.html

def calcGiniIndex(dataSet):
    """
    输入：数据集
    输出：数据集的基尼指数
    描述：计算给定数据集的基尼指数
    """
    numEntries = len(dataSet)                     # 返回数据集的行数
    labelCounts = {}                              # 保存每个标签（Label）出现次数的字典
    for featVec in dataSet:                       # 对每组特征向量进行统计
        currentLabel = featVec[-1]                  # 提取标签信息
        if currentLabel not in labelCounts.keys():    # 如果标签没有放入统计次数的字典，添加进去
            labelCounts[currentLabel] = 0
        labelCounts[currentLabel] += 1            # Label计数
    giniIndexEnt = 0.0                            # 基尼指数
    for key in labelCounts:                          # 计算基尼指数
        prob = float(labelCounts[key])/numEntries       # 选择该标签（Label）的概率
        giniIndexEnt += prob * (1.0 - prob)       # 利用公式计算
    return giniIndexEnt                              # 返回基尼指数

def splitDataSet(dataSet, axis, value):       # 待划分数据集合，特征下标，特征值
    """
    输入：数据集，选择维度，选择值
    输出：划分数据集
    描述：按照给定特征划分数据集；去除选择维度中等于选择值的项
    """
    retDataSet = []                          # 保存划分的数据子集
    for featVec in dataSet:                  # 遍历数据集中的每个样本
        if featVec[axis] == value:           #如果特征值符合要求，则添加到子集中
            reduceFeatVec = featVec[:axis]   # 保存第0到第axis-1个特征
            reduceFeatVec.extend(featVec[axis+1:])     # 保存第axis+1到最后一个特征
            retDataSet.append(reduceFeatVec)           # 添加符合要求的样本到划分子集中
    return retDataSet                                  # 返回划分好的（特征axis的值=value）的子集

def chooseBestFeatureToSplit(dataSet):
    """
    输入：数据集
    输出：最好的划分维度
    描述：选择最好的数据集划分维度
    """
    numFeatures = len(dataSet[0]) - 1                 # 特征数量
    bestInfoGini = calcGiniIndex(dataSet)             # 计算数据集的基尼指数
    bestFeature = -1                                   # 最优特征索引值
    for i in range(numFeatures):                        # 遍历所有特征
        featList = [example[i] for example in dataSet]      # 获取dataSet的第i个所有特征-第i列全部特征
        uniqueVals = set(featList)                    # 创建set集合{}元素不可重复
        newGini = 0.0
        for value in uniqueVals:                            # 计算新的基尼指数
            subDataSet = splitDataSet(dataSet, i, value)    # subDataSet划分后的子集
            prob = len(subDataSet)/float(len(dataSet))      # 计算子集概率
            newGini += prob * calcGiniIndex(subDataSet)     # 新的基尼指数
        if (newGini < bestInfoGini):
            bestInfoGini = newGini
            bestFeature = i
    return bestFeature                                         # 返回基尼指数最小的特征索引值

def majorityCnt(classList):
    """
    输入：分类类别列表
    输出：子节点的分类
    描述：数据集已经处理了所有属性，但是类标签依然不是唯一的，
          采用多数判决的方法决定该子节点的分类
    """
    classCount = {}
    for vote in classList:                # 统计classList中每个元素出现的次数
        if vote not in classCount.keys():
            classCount[vote] = 0
        classCount[vote] += 1
    sortedClassCount = sorted(classCount.iteritems(), key=operator.itemgetter(1), reversed=True)  # 根据字典的值降序排序
    return sortedClassCount[0][0]            # 返回classList中出现次数最多的元素

def createTree(dataSet, labels):
    """
    输入：数据集，特征标签
    输出：决策树
    描述：递归构建决策树，利用上述的函数
    """
    classList = [example[-1] for example in dataSet]
    if classList.count(classList[0]) == len(classList):
        # 类别完全相同，停止划分
        return classList[0]
    if len(dataSet[0]) == 1:
        # 遍历完所有特征时返回出现次数最多的
        return majorityCnt(classList)
    bestFeat = chooseBestFeatureToSplit(dataSet)
    bestFeatLabel = labels[bestFeat]
    myTree = {bestFeatLabel:{}}
    del(labels[bestFeat])
    # 得到列表包括节点所有的属性值
    featValues = [example[bestFeat] for example in dataSet]
    uniqueVals = set(featValues)
    for value in uniqueVals:
        subLabels = labels[:]
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), subLabels)
    return myTree

def classify(inputTree, featLabels, testVec):
    """
    输入：决策树，分类标签，测试数据
    输出：决策结果
    描述：跑决策树
    """
    firstStr = list(inputTree.keys())[0]
    secondDict = inputTree[firstStr]
    featIndex = featLabels.index(firstStr)
    for key in secondDict.keys():
        if testVec[featIndex] == key:
            if type(secondDict[key]).__name__ == 'dict':
                classLabel = classify(secondDict[key], featLabels, testVec)
            else:
                classLabel = secondDict[key]
    return classLabel

def classifyAll(inputTree, featLabels, testDataSet):
    """
    输入：决策树，分类标签，测试数据集
    输出：决策结果
    描述：跑决策树
    """
    classLabelAll = []
    for testVec in testDataSet:
        classLabelAll.append(classify(inputTree, featLabels, testVec))
    return classLabelAll

def storeTree(inputTree, filename):
    """
    输入：决策树，保存文件路径
    输出：
    描述：保存决策树到文件
    """
    import pickle
    fw = open(filename, 'wb')
    pickle.dump(inputTree, fw)
    fw.close()

def grabTree(filename):
    """
    输入：文件路径名
    输出：决策树
    描述：从文件读取决策树
    """
    import pickle
    fr = open(filename, 'rb')
    return pickle.load(fr)

def createDataSet():
    """
    outlook->  0: sunny | 1: overcast | 2: rain
    temperature-> 0: hot | 1: mild | 2: cool
    humidity-> 0: high | 1: normal
    windy-> 0: false | 1: true
    """
    dataSet = [[0, 0, 0, 0, 'N'],
               [0, 0, 0, 1, 'N'],
               [1, 0, 0, 0, 'Y'],
               [2, 1, 0, 0, 'Y'],
               [2, 2, 1, 0, 'Y'],
               [2, 2, 1, 1, 'N'],
               [1, 2, 1, 1, 'Y']]
    labels = ['outlook', 'temperature', 'humidity', 'windy']
    return dataSet, labels

def createTestSet():
    """
    outlook->  0: sunny | 1: overcast | 2: rain
    temperature-> 0: hot | 1: mild | 2: cool
    humidity-> 0: high | 1: normal
    windy-> 0: false | 1: true
    """
    testSet = [[0, 1, 0, 0],
               [0, 2, 1, 0],
               [2, 1, 1, 0],
               [0, 1, 1, 1],
               [1, 1, 0, 1],
               [1, 0, 1, 0],
               [2, 1, 0, 1]]
    return testSet

def main():
    dataSet, labels = createDataSet()
    labels_tmp = labels[:] # 拷贝，createTree会改变labels
    desicionTree = createTree(dataSet, labels_tmp)
    #storeTree(desicionTree, 'classifierStorage.txt')
    #desicionTree = grabTree('classifierStorage.txt')
    print('desicionTree:\n', desicionTree)
    testSet = createTestSet()
    print('classifyResult:\n', classifyAll(desicionTree, labels, testSet))

if __name__ == '__main__':
    main()